package dslab.util;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;


public class cryptoutil {

    public static void main(String[] args) throws NoSuchAlgorithmException, IOException, InvalidKeySpecException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
    /*
        Path pubPath = Paths.get("/home/caritas/dslab/keys/client/transfer-1_pub.der");
        Path privPath = Paths.get("/home/caritas/dslab/keys/server/transfer-1.der");
        PublicKey publicKey = getPublicKey(pubPath);
        PrivateKey privateKey = getPrivateKey(privPath);
        String message = "Ich mag Schnitzel :3";
        System.out.println(message);

        Cipher encCipher = createRSAEncryptCipher(publicKey);
        System.out.println(encCipher.doFinal(message.getBytes(StandardCharsets.UTF_8)).length);
        String encMessage64 = new String(Base64.getEncoder().encode(encCipher.doFinal(message.getBytes(StandardCharsets.UTF_8))));
        System.out.println(encMessage64.length());
        System.out.println(encMessage64);

        Cipher decCipher = createRSADecryptCipher(privateKey);
        String decMessage = new String(decCipher.doFinal(Base64.getDecoder().decode(encMessage64.getBytes(StandardCharsets.UTF_8))));
        System.out.println(decMessage);

     */
        String message = "Fuck Java";
        Cipher cipher = Cipher.getInstance("AES/CTR/NoPadding");
        Pair<SecretKey, byte[]> keyData = generateNewKeydata();
        cipher.init(Cipher.ENCRYPT_MODE, keyData.key(), new IvParameterSpec(keyData.value()));
        String enc = encryptMessage(message, cipher);
        String key = encodeB64(keyData.key().getEncoded());
        String iv = encodeB64(keyData.value());


        SecretKey secret = decodeSecret64(key);
        cipher.init(Cipher.DECRYPT_MODE, secret, new IvParameterSpec(decodeB64(iv)));
        String messdec = decryptMessage(enc, cipher);
        System.out.println(messdec);


    }



    private static SecretKey genSecret() throws NoSuchAlgorithmException {
        KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
        keyGenerator.init(256);
        return keyGenerator.generateKey();
    }

    private static byte[] genRandom(int size) {
        byte[] rand = new byte[size];
        new SecureRandom().nextBytes(rand);
        return rand;
    }

    private static Pair<SecretKey, byte[]> generateNewKeydata() throws NoSuchAlgorithmException {
        return new Pair<>(genSecret(), genRandom(16));
    }

    private static SecretKey decodeSecret64(String key) {
        return new SecretKeySpec(decodeB64(key), 0, decodeB64(key).length, "AES");
    }

    private static String encodeB64(String message) {
        return Base64.getEncoder().encodeToString(message.getBytes(StandardCharsets.UTF_8));
    }

    private static byte[] decodeB64(String message) {
        return Base64.getDecoder().decode(message);
    }

    private static String encodeB64(byte[] message) {
        return Base64.getEncoder().encodeToString(message);
    }

    private static byte[] decodeB64(byte[] message) {
        return Base64.getDecoder().decode(message);
    }

    private static PrivateKey getPrivateKey(Path location) {
        try {
            PKCS8EncodedKeySpec keySpecPriv = new PKCS8EncodedKeySpec(Files.readAllBytes(location));
            return KeyFactory.getInstance("RSA").generatePrivate(keySpecPriv);
        } catch (Exception exp) {
            System.err.println(exp.getMessage());
            return null;
        }
    }

    private static PublicKey getPublicKey(Path location) {
        try {
            X509EncodedKeySpec keySpecPub = new X509EncodedKeySpec(Files.readAllBytes(location));
            return KeyFactory.getInstance("RSA").generatePublic(keySpecPub);
        } catch (Exception exp) {
            System.err.println(exp.getMessage());
            return null;
        }
    }

    private static Cipher createRSAEncryptCipher(PublicKey key) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.ENCRYPT_MODE, key);
            return cipher;
        } catch (Exception e) {
            System.err.println(e.getMessage());
            return null;
        }
    }

    private static Cipher createRSADecryptCipher(PrivateKey key) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.DECRYPT_MODE, key);
            return cipher;
        } catch (Exception e) {
            System.err.println(e.getMessage());
            return null;
        }
    }

    private static String encryptMessage(String message, Cipher cipher) throws IllegalBlockSizeException, BadPaddingException {
        return new String(Base64.getEncoder().encode(cipher.doFinal(message.getBytes(StandardCharsets.UTF_8))));
    }

    private static String decryptMessage(String message, Cipher cipher) throws IllegalBlockSizeException, BadPaddingException {
        return new String(cipher.doFinal(Base64.getDecoder().decode(message.getBytes(StandardCharsets.UTF_8))));
    }
}
